#include "vec_add.h"

__global__ void cuda_vec_add(float *out, float *x, float *y, size_t N) {
  size_t thread_grid_idx = size_t(blockIdx.x * blockDim.x + threadIdx.x);
  size_t num_threads_in_grid = size_t(blockDim.x * gridDim.x);
  for (size_t i = thread_grid_idx; i < N; i += num_threads_in_grid) {
    out[i] = x[i] + y[i];
  }
}

void vec_add(float *out, float *x, float *y, cudaStream_t stream, size_t N) {
  constexpr std::size_t threads_per_block = 512;
  auto num_blocks = (N + threads_per_block - 1) / threads_per_block;
  cuda_vec_add<<<num_blocks, threads_per_block, 0, stream>>>(out, x, y, N);
}